/* Janus Chat - Version 1.0
 * Authors: Matthew Perry
 *          Michael Freid
 *          Marc Latou
 *          Richard Lee
 *
 * The Janus Chat Client is a highly secured
 * peer-to-peer chat client designed to facilitate
 * the transfer of sensitive data or as a communication
 * medium for the extra paranoid
 *
 * This file contains the header for a wrapper
 * of the crypto++ RSA library so that Janus can encrypt
 * its session keys for AES message encryption */

#include "jAes.h"
#include "exceptions.h"

#include <cryptopp/filters.h>
using CryptoPP::StringSink;
using CryptoPP::StringSource;
using CryptoPP::AuthenticatedEncryptionFilter;
using CryptoPP::AuthenticatedDecryptionFilter;

#include <cryptopp/osrng.h>
using CryptoPP::AutoSeededRandomPool;

#include <cryptopp/aes.h>
using CryptoPP::AES;

#include <cryptopp/gcm.h>
using CryptoPP::GCM;

using std::string;

/* Static serialize function:
 * this is a helper function used by the jAes
 * class to take a struct representation of the
 * components of a session key and encode them into
 * a sendable string */
string jAes::serialize( struct sessionKey* key )
{
    // Create a character array of appropriate size
    char encoding[ 8 + key->keyLen + key->ivLen ];    

    // Set first 4 bytes to keyLen
    for( int i = 3; i >= 0; --i )
        encoding[i] = char( key->keyLen >> ( ( 3 - i ) * 8 ) );

    // Set second 4 bytes to ivLen
    for( int i = 7; i >= 4; --i )
        encoding[i] = char( key->ivLen >> ( ( 7 - i ) * 8 ) );

    // Set remaining bytes to key + iv
    for( unsigned int i = 0; i < key->keyLen; ++i )
        encoding[i + 8] = key->key[i];
    for( unsigned int i = 0; i < key->ivLen; ++i )
        encoding[i + key->keyLen + 8] = key->iv[i];

    string retval( encoding, key->keyLen + key->ivLen + 8 );
    return retval;
}

/* Static revover function:
 * this is a helper function used by the jAes
 * class to take a string representation of the
 * components of a session key and decode them into
 * a usable struct */
struct jAes::sessionKey jAes::recover( const string& sKey )
{
    // Create the empty struct to return
    struct jAes::sessionKey retKey;

    // Extract the header portion from the serial key
    string header = sKey.substr( 0, 8 );
    // Set the keyLen to the first 4 bytes of the header
    retKey.keyLen = 0;
    for( int i = 0; i < 4; ++i )
        retKey.keyLen = (retKey.keyLen << 8) | (unsigned char)header[i];

    // Set the ivLen to the second 4 bytes of the header
    retKey.ivLen = 0;
    for( int i = 4; i < 8; ++i )
        retKey.ivLen = (retKey.ivLen << 8) | (unsigned char)header[i];

    // Extract the key data from the serial key
    string keyData = sKey.substr( 8, retKey.keyLen );
    // Allocate memory for key
    retKey.key = new unsigned char[retKey.keyLen];
    for( unsigned int i = 0; i < retKey.keyLen; ++i )
        retKey.key[i] = keyData[i];

    // Extract the iv data from the serial key
    string ivData = sKey.substr( 8 + retKey.keyLen, retKey.ivLen );
    // Allocate memory for iv
    retKey.iv = new unsigned char[retKey.ivLen];
    for( unsigned int i = 0; i < retKey.ivLen; ++i )
        retKey.iv[i] = ivData[i];

    return retKey;
}

/* Static generate key function:
 * This function is used to give the user a
 * session key to use with the AES encryption
 * scheme in this class */
string jAes::generateKey( int size )
{
    // Ensure that the size is proper
    if( size != 16 && size != 24 && size != 32 )
        throw new MessageException( "Incorrect key size" );

    // Create the empty struct
    struct jAes::sessionKey genKey;

    // Create a random number generator
    AutoSeededRandomPool rnd;

    // Generate a random key
    unsigned char key[size];
    rnd.GenerateBlock( key, size );
    
    // Generate a random iv
    unsigned char iv[AES::BLOCKSIZE];
    rnd.GenerateBlock( iv, AES::BLOCKSIZE );

    // Set the return values
    genKey.key = key;
    genKey.keyLen = size;
    genKey.iv = iv;
    genKey.ivLen = AES::BLOCKSIZE;

    // Return a serial string representation
    return serialize( &genKey );
}

/* Static encrypt function:
 * Given a serial session key, sKey and a
 * plaintext string this function will encrypt
 * the string using the AES_GCM method */
string jAes::encrypt( const string& sKey, const string& plaintext )
{
    // Create cipher string placeholder and recover the key from serial form
    string cipher;
    struct jAes::sessionKey key = recover( sKey );

    // Create the encryption with proper keys
    GCM<AES>::Encryption enc;
    enc.SetKeyWithIV( key.key, key.keyLen, key.iv, key.ivLen );

    // Set up authentication filter
    AuthenticatedEncryptionFilter aef( enc, new StringSink( cipher ) );

    // Encrypt plain text
    aef.Put( (const unsigned char*)plaintext.data(), plaintext.size() );
    aef.MessageEnd();

    return cipher;
}

/* Static decrypt function:
 * Given a serial session key, sKey and a
 * cipher string this function will decrypt
 * the string using the AES_GCM method */
string jAes::decrypt( const string& sKey, const string& cipher )
{
    // Create plaintext placeholder and recover the key from serial form
    string plaintext;
    struct jAes::sessionKey key = recover( sKey );

    //Create decryption with proper keys
    GCM<AES>::Decryption dec;
    dec.SetKeyWithIV( key.key, key.keyLen, key.iv, key.ivLen );

    // Set up authentication filter
    AuthenticatedDecryptionFilter adf( dec, new StringSink( plaintext ) );

    // Decrypt cipher text
    adf.Put( (const unsigned char*)cipher.data(), cipher.size() );
    adf.MessageEnd();

    return plaintext;
}
